import pandas
import numpy


def add_rets(df, pricecol, signcol='sign', possigncol='sign', fwds=None):
  if fwds is None:
    fwds = [15, 30]

  pos_mtmp = None
  neg_mtmp = None
  pos_mtmps = []
  neg_mtmps = []

  inv_signs = df['sign'][::-1]
  inv_prcs = df[pricecol][::-1]

  for inv_sign, inv_prc in zip(inv_signs, inv_prcs):
    if inv_sign > 0:
      pos_mtmp = inv_prc
    else:
      neg_mtmp = inv_prc
    pos_mtmps.append(pos_mtmp)
    neg_mtmps.append(neg_mtmp)

  df['pos_mtmp'] = pos_mtmps[::-1]
  df['neg_mtmp'] = neg_mtmps[::-1]

  df['mtmp'] = 0.5 * (df['pos_mtmp'] + df['neg_mtmp'])

  next_mtmp = None
  next_mtmps = []
  next_ts = None
  rows = []

  inv_mtmps = df['mtmp'][::-1]
  inv_poss = df[possigncol][::-1]
  inv_tss = df['timestamp'][::-1]

  for inv_mtmp, inv_pos, inv_ts in zip(inv_mtmps, inv_poss, inv_tss):

    row = {}
    new_mtmp = inv_mtmp
    ret = 0 if next_mtmp is None else (next_mtmp - new_mtmp) * inv_pos

    for fwd in fwds:
      ret_fwd = 0 if next_mtmp is None else (numpy.mean(next_mtmps[-fwd:]) - new_mtmp) * inv_pos
      row['next_ret%s' % fwd] = ret_fwd

    if next_ts is not None and next_ts - inv_ts > 600 * 1e9:
      ret = 0

    row['next_ret'] = ret
    next_mtmp = new_mtmp
    next_ts = inv_ts
    next_mtmps.append(new_mtmp)
    rows.append(row)
  df = df.reset_index(drop=True)
  df2 = pandas.DataFrame(rows[::-1])
  return pandas.concat([df, df2], axis=1)


def add_signed_rets(df, pricecol, possigncol='pos_sign'):
  df['neg_mtmp'] = df[pricecol].copy()
  df.loc[df['sign'] > 0, 'neg_mtmp'] = numpy.nan
  mtmp = df['neg_mtmp'].fillna(method='backfill')
  ret_neg = numpy.nan_to_num(mtmp.diff(-1))

  df['pos_mtmp'] = df[pricecol].copy()
  df.loc[df['sign'] <= 0, 'pos_mtmp'] = numpy.nan
  mtmp = df['pos_mtmp'].fillna(method='backfill')
  ret_pos = numpy.nan_to_num(mtmp.diff(-1))

  ret = ret_neg.copy()
  ret[df['sign'] > 0] = ret_pos[df['sign'] > 0]

  df['next_ret_abs'] = ret
  df['next_ret'] = ret * -df[possigncol]
  return df


def add_unsigned_rets(df, pricecol, possigncol='pos_sign'):
  df['mtmp'] = df[pricecol].copy()
  mtmp = df['mtmp'].fillna(method='backfill')
  ret = numpy.nan_to_num(mtmp.diff(-1))
  df['next_ret_abs'] = ret
  df['next_ret'] = ret * -df[possigncol]
  return df
